import pandas as pd
import os
from io import StringIO
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick

path = r'C:\Users\Jeff Levy\Desktop\replication'

def parse_bepuc3(fname):
    df = pd.read_stata(os.path.join(path, fname))
    cols = {'step':'step', 'D03':'IRF_epuc', 'pD03_ll':'Lower_epuc', 'pD03_ul':'Upper_epuc',
            'D06':'IRF_bepuc', 'pD06_ll':'Lower_bepuc', 'pD06_ul':'Upper_bepuc',
            'D04':'IRF_bono', 'pD04_ll':'Lower_bono', 'pD04_ul':'Upper_bono',
            'FEVD03':'FEVD_epuc', 'FEVD06':'FEVD_bepuc', 'FEVD09':'FEVD_unemp'}
    df = df.rename(columns=cols)
    return df[list(cols.values())]

def irf_plots_bepuc(data):
    #manually calculated from xtvar_ready.csv, 44 state sample (using modified xtvar.do from Dropbox)
    sdr = {'bepuc3_uf_bepuc2nd.dta':.44233373,
           'bepuc3_hp_bepuc2nd.dta':.43125623,
           'bepuc3_hf_bepuc2nd.dta':.44419685}
    sds = {'bepuc3_uf_bepuc2nd.dta':.28194746,
           'bepuc3_hp_bepuc2nd.dta':.27522176,
           'bepuc3_hf_bepuc2nd.dta':.28636243}

    for figure in ['epuc', 'bepuc']:
        fig, axs = plt.subplots(3, 1, figsize=(8,10), sharey=True)
        for kind, ax in zip(['hf', 'hp', 'uf'], axs):
            fname = f'bepuc3_{kind}_bepuc2nd.dta'
            sd_resid = sdr[fname]
            sd_struc = sds[fname]
            d = data[fname]

            irf, lower, upper = f'IRF_{figure}', f'Lower_{figure}', f'Upper_{figure}'
            x = d['step']
            ax.plot(x, d[irf], 'k-')
            ax.plot(x, d[lower], 'k--', alpha=0.5)
            ax.plot(x, d[upper], 'k--', alpha=0.5)
            ax.axhline(0, 0, 1, color='k', linestyle='-', linewidth=0.2)

            peak = d[irf].max()
            peak_pos = peak*1.05
            peak_date = d[d[irf] == peak].step.values[0]

            ax.annotate(str(round(peak, 3)).ljust(5, "0"), (peak_date, peak_pos))

            if figure == 'epuc':
                e_value = str(round(sd_resid, 3)).ljust(5, "0")
                unit_sd_size = f'EPU-C \u03C3: {e_value}'#\n{e}: {e_value}'
                ax.text(.98, .91, unit_sd_size.ljust(3, '0'), fontsize='large', horizontalalignment='right', transform=ax.transAxes)
            elif figure == 'bepuc':
                e_value = str(round(sd_struc, 3)).ljust(5, "0")
                unit_sd_size = f'BEPU-C \u03C3: {e_value}'#\n{e}: {e_value}'
                ax.text(.98, .91, unit_sd_size.ljust(3, '0'), fontsize='large', horizontalalignment='right', transform=ax.transAxes)

        axs[0].set_ylabel('Hamilton Filter', fontsize='xx-large')
        axs[1].set_ylabel('HP Filter', fontsize='xx-large')
        axs[2].set_ylabel('No Filter', fontsize='xx-large')
        for a in axs:
            a.xaxis.set_tick_params(labelsize='x-large')
            a.yaxis.set_tick_params(labelsize='x-large')

        out = 'bepuc3_' + f'imp{figure.capitalize()}_' + '.png'
        fig.tight_layout()
        fig.savefig(os.path.join(path, 'irfs', out))
        print('Saved:', os.path.join(path, 'irfs', out))
        plt.close()

fnames = [f for f in os.listdir(path) if f.endswith('.dta') and 'bepuc3' in f]
data = {f:parse_bepuc3(f) for f in fnames}
irf_plots_bepuc(data)
